import os
import argparse
import contextlib

source_domains = ["ctrl", "fair_wmt19", "gpt2_xl", "gpt3", "grover_mega", "xlm", "chatgpt"]
target_domains = ["ctrl", "fair_wmt19", "gpt2_xl", "gpt3", "grover_mega", "xlm", "chatgpt"]

loss_type = 'simclr'

for src in source_domains:
    for tgt in target_domains:
        if src != tgt:
            parser = argparse.ArgumentParser()

            parser.add_argument('--max-epochs', type=int, default=None,
                                help='最大训练周期数。')
            parser.add_argument('--device', type=str, default=None,
                                help='用于训练的设备，例如 "cuda" 或 "cpu"。')
            parser.add_argument('--batch-size', type=int, default=16,
                                help='训练批次大小。')
            parser.add_argument('--max-sequence-length', type=int, default=256,
                                help='序列的最大长度。')
            parser.add_argument('--random-sequence-length', action='store_true',
                                help='是否随机选择序列长度。')
            parser.add_argument('--epoch-size', type=int, default=None,
                                help='每个周期的数据量大小。')
            parser.add_argument('--loss_type', type=str, default=loss_type,
                                help='损失函数类型，可选 [simclr, supcon]。')
            parser.add_argument('--seed', type=int, default=None,
                                help='随机种子，用于结果的复现。')

            # 源域数据路径和名称
            if src == 'chatgpt':
                parser.add_argument('--src_data-dir', type=str, default='/home/abhatt43/projects/Data_for_Testing/syn_rep/ChatGPT/',
                                    help='源域数据目录。')
                parser.add_argument('--src_real-dataset', type=str, default=src+'_real',
                                    help='源域真实数据集名称。')
                parser.add_argument('--src_fake-dataset', type=str, default=src+'_fake',
                                    help='源域合成数据集名称。')
            else:
                parser.add_argument('--src_data-dir', type=str, default='/home/abhatt43/projects/Data_for_Testing/syn_rep/TuringBench/TT_'+src+'/',
                                    help='源域数据目录。')
                parser.add_argument('--src_real-dataset', type=str, default='tb_tt_'+src+'_real',
                                    help='源域真实数据集名称。')
                parser.add_argument('--src_fake-dataset', type=str, default='tb_tt_'+src+'_fake',
                                    help='源域合成数据集名称。')

            # 目标域数据路径和名称
            if tgt == 'chatgpt':
                parser.add_argument('--tgt_data-dir', type=str, default='/home/abhatt43/projects/Data_for_Testing/syn_rep/ChatGPT/',
                                    help='目标域数据目录。')
                parser.add_argument('--tgt_real-dataset', type=str, default=tgt+'_real',
                                    help='目标域真实数据集名称。')
                parser.add_argument('--tgt_fake-dataset', type=str, default=tgt+'_fake',
                                    help='目标域合成数据集名称。')
            else:
                parser.add_argument('--tgt_data-dir', type=str, default='/home/abhatt43/projects/Data_for_Testing/syn_rep/TuringBench/TT_'+tgt+'/',
                                    help='目标域数据目录。')
                parser.add_argument('--tgt_real-dataset', type=str, default='tb_tt_'+tgt+'_real',
                                    help='目标域真实数据集名称。')
                parser.add_argument('--tgt_fake-dataset', type=str, default='tb_tt_'+tgt+'_fake',
                                    help='目标域合成数据集名称。')

            parser.add_argument('--model_save_path', default=os.getcwd()+'/models/',
                                help='模型保存路径。')

            if loss_type == 'simclr':
                parser.add_argument('--model_save_name', default=src+'_'+tgt+'_syn_rep_'+loss_type+'.pt',
                                    help='模型保存名称。')

            parser.add_argument('--token-dropout', type=float, default=None,
                                help='令牌丢弃率。')

            parser.add_argument('--large', action='store_true',
                                help='使用roberta-large模型而非roberta-base。')
            parser.add_argument('--learning-rate', type=float, default=2e-5,
                                help='学习率。')
            parser.add_argument('--weight-decay', type=float, default=0,
                                help='权重衰减。')
            parser.add_argument('--load-decay', type=float, default=0,
                                help='加载衰减。')
            parser.add_argument('--lambda_w', type=float, default=0.5,
                                help='权重因子λ。')

            args = parser.parse_args(args=['--max-epochs=5', '--model_save_path='+os.getcwd()+'/models/'])

            filename = 'output_log_'+src+'_'+tgt+'_syn_rep_'+loss_type+'.log'
            filepath = os.getcwd() + '/output_logs/'
            with open(os.path.join(filepath, filename), 'w') as f:
                with contextlib.redirect_stderr(f):
                    main(args)
